
package org.apache.solr.cloud;

import java.io.File;
import java.io.IOException;
import java.util.List;
import java.util.concurrent.TimeoutException;
import javax.xml.parsers.ParserConfigurationException;
import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.CommandLineParser;
import org.apache.commons.cli.HelpFormatter;
import org.apache.commons.cli.Option;
import org.apache.commons.cli.OptionBuilder;
import org.apache.commons.cli.Options;
import org.apache.commons.cli.ParseException;
import org.apache.commons.cli.PosixParser;
import org.apache.solr.common.cloud.OnReconnect;
import org.apache.solr.common.cloud.SolrZkClient;
import org.apache.solr.core.Config;
import org.apache.solr.core.SolrResourceLoader;
import org.apache.zookeeper.KeeperException;
import org.xml.sax.InputSource;
import org.xml.sax.SAXException;

public class ZkCLI {

    private static final String MAKEPATH    = "makepath";
    private static final String DOWNCONFIG  = "downconfig";
    private static final String ZK_CLI_NAME = "ZkCLI";
    private static final String HELP        = "help";
    private static final String LINKCONFIG  = "linkconfig";
    private static final String CONFDIR     = "confdir";
    private static final String CONFNAME    = "confname";
    private static final String ZKHOST      = "zkhost";
    private static final String RUNZK       = "runzk";
    private static final String SOLRHOME    = "solrhome";
    private static final String BOOTSTRAP   = "bootstrap";
    private static final String SOLR_XML    = "solr.xml";
    private static final String UPCONFIG    = "upconfig";
    private static final String COLLECTION  = "collection";
    private static final String CLEAR       = "clear";
    private static final String CMD         = "cmd";

    /**
     * Allows you to perform a variety of zookeeper related tasks, such as:
     *
     * Bootstrap the current configs for all collections in solr.xml.
     *
     * Upload a named config set from a given directory.
     *
     * Link a named config set explicity to a collection.
     *
     * Clear ZooKeeper info.
     *
     * If you also pass a solrPort, it will be used to start an embedded zk useful
     * for single machine, multi node tests.
     */
    public static void main(String[] args) throws InterruptedException, TimeoutException, IOException, ParserConfigurationException, SAXException, KeeperException {

        CommandLineParser parser = new PosixParser();
        Options options = new Options();

        options.addOption(OptionBuilder.hasArg(true).withDescription(
                "cmd to run: " + BOOTSTRAP + ", " + UPCONFIG + ", " + DOWNCONFIG
                + ", " + LINKCONFIG + ", " + MAKEPATH + ", " + CLEAR).create(CMD));

        Option zkHostOption = new Option("z", ZKHOST, true, "ZooKeeper host address");
        options.addOption(zkHostOption);
        Option solrHomeOption = new Option("s", SOLRHOME, true, "for " + BOOTSTRAP + ", " + RUNZK + ": solrhome location");
        options.addOption(zkHostOption);
        options.addOption(solrHomeOption);

        options.addOption("d", CONFDIR, true, "for " + UPCONFIG + ": a directory of configuration files");
        options.addOption("n", CONFNAME, true, "for " + UPCONFIG + ", " + LINKCONFIG + ": name of the config set");
        options.addOption("c", COLLECTION, true, "for " + LINKCONFIG + ": name of the collection");
        options.addOption("r", RUNZK, true, "run zk internally by passing the solr run port - only for clusters on one machine (tests, dev)");
        options.addOption("h", HELP, false, "bring up this help page");

        try {
            // parse the command line arguments
            CommandLine line = parser.parse(options, args);

            if(line.hasOption(HELP) || !line.hasOption(ZKHOST) || !line.hasOption(CMD)) {
                // automatically generate the help statement
                HelpFormatter formatter = new HelpFormatter();
                formatter.printHelp(ZK_CLI_NAME, options);
                System.out.println("Examples:");
                System.out.println("zkcli.sh -cmd " + BOOTSTRAP + " -" + SOLRHOME + " /opt/solr");
                System.out.println("zkcli.sh -cmd " + UPCONFIG + " -" + CONFDIR + " /opt/solr/collection1/conf" + " -" + CONFNAME + " myconf");
                System.out.println("zkcli.sh -cmd " + DOWNCONFIG + " -" + CONFDIR + " /opt/solr/collection1/conf" + " -" + CONFNAME + " myconf");
                System.out.println("zkcli.sh -cmd " + LINKCONFIG + " -" + COLLECTION + " collection1" + " -" + CONFNAME + " myconf");
                System.out.println("zkcli.sh -cmd " + MAKEPATH + " /apache/solr");
                System.out.println("zkcli.sh -cmd " + CLEAR + " /solr");
                return;
            }

            // start up a tmp zk server first
            String zkServerAddress = line.getOptionValue(ZKHOST);
            String solrHome = line.getOptionValue(SOLRHOME);

            String solrPort = null;
            if(line.hasOption(RUNZK)) {
                if(!line.hasOption(SOLRHOME)) {
                    System.out
                            .println("-" + SOLRHOME + " is required for " + RUNZK);
                    System.exit(1);
                }
                solrPort = line.getOptionValue(RUNZK);
            }

            SolrZkServer zkServer = null;
            if(solrPort != null) {
                zkServer = new SolrZkServer("true", null, solrHome + "/zoo_data", solrHome, solrPort);
                zkServer.parseConfig();
                zkServer.start();
            }
            SolrZkClient zkClient = null;
            try {
                zkClient = new SolrZkClient(zkServerAddress, 30000, 30000,
                        new OnReconnect() {
                            @Override
                            public void command() {
                            }
                        });

                if(line.getOptionValue(CMD).equals(BOOTSTRAP)) {
                    if(!line.hasOption(SOLRHOME)) {
                        System.out.println("-" + SOLRHOME + " is required for " + BOOTSTRAP);
                        System.exit(1);
                    }
                    SolrResourceLoader loader = new SolrResourceLoader(solrHome);
                    solrHome = loader.getInstanceDir();

                    InputSource cfgis = new InputSource(new File(solrHome, SOLR_XML).toURI().toASCIIString());
                    Config cfg = new Config(loader, null, cfgis, null, false);
                    ZkController.bootstrapConf(zkClient, cfg, solrHome);
                }
                else if(line.getOptionValue(CMD).equals(UPCONFIG)) {
                    if(!line.hasOption(CONFDIR) || !line.hasOption(CONFNAME)) {
                        System.out.println("-" + CONFDIR + " and -" + CONFNAME + " are required for " + UPCONFIG);
                        System.exit(1);
                    }
                    String confDir = line.getOptionValue(CONFDIR);
                    String confName = line.getOptionValue(CONFNAME);

                    ZkController.uploadConfigDir(zkClient, new File(confDir), confName);
                }
                else if(line.getOptionValue(CMD).equals(DOWNCONFIG)) {
                    if(!line.hasOption(CONFDIR) || !line.hasOption(CONFNAME)) {
                        System.out.println("-" + CONFDIR + " and -" + CONFNAME
                                + " are required for " + DOWNCONFIG);
                        System.exit(1);
                    }
                    String confDir = line.getOptionValue(CONFDIR);
                    String confName = line.getOptionValue(CONFNAME);

                    ZkController.downloadConfigDir(zkClient, confName, new File(confDir));
                }
                else if(line.getOptionValue(CMD).equals(LINKCONFIG)) {
                    if(!line.hasOption(COLLECTION) || !line.hasOption(CONFNAME)) {
                        System.out.println("-" + CONFDIR + " and -" + CONFNAME + " are required for " + LINKCONFIG);
                        System.exit(1);
                    }
                    String collection = line.getOptionValue(COLLECTION);
                    String confName = line.getOptionValue(CONFNAME);

                    ZkController.linkConfSet(zkClient, collection, confName);
                }
                else if(line.getOptionValue(CMD).equals(CLEAR)) {
                    List arglist = line.getArgList();
                    if(arglist.size() != 1) {
                        System.out.println("-" + CLEAR + " requires one arg - the path to clear");
                        System.exit(1);
                    }
                    zkClient.clean(arglist.get(0).toString());
                }
                else if(line.getOptionValue(CMD).equals(MAKEPATH)) {
                    List arglist = line.getArgList();
                    if(arglist.size() != 1) {
                        System.out.println("-" + MAKEPATH + " requires one arg - the path to make");
                        System.exit(1);
                    }
                    zkClient.makePath(arglist.get(0).toString(), true);
                }
            }
            finally {
                if(solrPort != null) {
                    zkServer.stop();
                }
                if(zkClient != null) {
                    zkClient.close();
                }
            }
        }
        catch(ParseException exp) {
            System.out.println("Unexpected exception:" + exp.getMessage());
        }
    }
}
